import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, models
from sklearn.model_selection import train_test_split
from PIL import Image
import pandas as pd
import numpy as np
from tqdm import tqdm

# 设置随机种子保证可复现性
torch.manual_seed(42)
np.random.seed(42)

# 设备配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# ----------------------
# 1. 数据准备
# ----------------------
class FaceDataset(Dataset):
    def __init__(self, dataframe, root_dir, transform=None):
        self.dataframe = dataframe
        self.root_dir = root_dir
        self.transform = transform

    def __len__(self):
        return len(self.dataframe)

    def __getitem__(self, idx):
        img_path = os.path.join(self.root_dir, self.dataframe.iloc[idx]['image_path'])
        image = Image.open(img_path).convert('RGB')
        label = self.dataframe.iloc[idx]['label']
        
        if self.transform:
            image = self.transform(image)
            
        return image, label

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # ResNet50的标准输入尺寸
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # ImageNet标准化
])

# 从Kaggle数据集创建DataFrame
def create_dataframe(root_dir, datatype='train'):
    real_dir = os.path.join(root_dir, datatype, 'real')
    fake_dir = os.path.join(root_dir, datatype, 'fake')
    
    real_images = [{'image_path': os.path.join(datatype, 'real', f), 'label': 1} for f in os.listdir(real_dir)]
    fake_images = [{'image_path': os.path.join(datatype, 'fake', f), 'label': 0} for f in os.listdir(fake_dir)]
    
    return pd.DataFrame(real_images + fake_images)

# 创建数据集
root_dir = './140k-real-and-fake-faces/real_vs_fake/real-vs-fake'
train_df = create_dataframe(root_dir)
val_df = create_dataframe(root_dir, datatype='valid')

# 创建DataLoader
train_dataset = FaceDataset(train_df, root_dir, transform=transform)
val_dataset = FaceDataset(val_df, root_dir, transform=transform)

batch_size = 32
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

# ----------------------
# 2. 模型构建
# ----------------------
def build_model(pretrained=True):
    model = models.resnet50(pretrained=False)
    
    # 冻结所有卷积层
    # for param in model.parameters():
    #     param.requires_grad = False
    
    # 替换最后的全连接层
    num_features = model.fc.in_features
    model.fc = nn.Sequential(
        nn.Linear(num_features, 256),
        nn.ReLU(),
        nn.Dropout(0.5),
        nn.Linear(256, 1),
        nn.Sigmoid()
    )
    if os.path.exists('best_model.pth'):
        print("Loading locally saved best model...")
        model.load_state_dict(torch.load('best_model.pth'))
        print("Successfully loaded local model weights!")
    elif pretrained:
        print("Loading ImageNet pretrained weights...")
        pretrained_model = models.resnet50(pretrained=True)
        
        # 复制除最后一层外的所有权重
        model_dict = model.state_dict()
        pretrained_dict = {k: v for k, v in pretrained_model.state_dict().items() 
                          if k in model_dict and not k.startswith('fc')}
        model_dict.update(pretrained_dict)
        model.load_state_dict(model_dict)
        print("Successfully loaded pretrained weights (except last layer)!")
    else:
        print("Training from scratch (no weights loaded)")
    return model.to(device)

model = build_model(pretrained=True)

# 定义损失函数和优化器
criterion = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# ----------------------
# 3. 训练和验证
# ----------------------
def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=10):
    best_acc = 0.0
    
    for epoch in range(num_epochs):
        print(f'Epoch {epoch+1}/{num_epochs}')
        print('-' * 10)
        
        # 训练阶段
        model.train()
        running_loss = 0.0
        running_corrects = 0
        
        for inputs, labels in tqdm(train_loader, desc='Training'):
            inputs = inputs.to(device)
            labels = labels.float().unsqueeze(1).to(device)
            
            optimizer.zero_grad()
            
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item() * inputs.size(0)
            preds = (outputs > 0.5).float()
            running_corrects += torch.sum(preds == labels.data)
            
        epoch_loss = running_loss / len(train_dataset)
        epoch_acc = running_corrects.double() / len(train_dataset)
        
        print(f'Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')
        
        # 验证阶段
        model.eval()
        val_loss = 0.0
        val_corrects = 0
        
        with torch.no_grad():
            for inputs, labels in tqdm(val_loader, desc='Validating'):
                inputs = inputs.to(device)
                labels = labels.float().unsqueeze(1).to(device)
                
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                
                val_loss += loss.item() * inputs.size(0)
                preds = (outputs > 0.5).float()
                val_corrects += torch.sum(preds == labels.data)
                
            val_loss = val_loss / len(val_dataset)
            val_acc = val_corrects.double() / len(val_dataset)
            
            print(f'Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}')
            
            # 保存最佳模型
            if val_acc > best_acc:
                best_acc = val_acc
                torch.save(model.state_dict(), 'best_model.pth')
                
        print()
    
    print(f'Best val Acc: {best_acc:.4f}')

# 开始训练
train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=5)

# ----------------------
# 4. 测试模型
# ----------------------
def test_model(model, test_loader):
    model.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs = inputs.to(device)
            labels = labels.float().unsqueeze(1).to(device)
            
            outputs = model(inputs)
            preds = (outputs > 0.5).float()
            
            total += labels.size(0)
            correct += (preds == labels).sum().item()
    
    print(f'Test Accuracy: {100 * correct / total:.2f}%')

# 测试
test_df = create_dataframe(root_dir, datatype='test')
test_dataset = FaceDataset(test_df, root_dir, transform=transform)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
test_model(model, test_loader)